#include "rbt.h"


/*--------------------------以下是一些辅助函数--------------------------*/
/**
 * @brief 获取结点的颜色
 * 
 * @param node 
 * @return int 
 */
int get_node_color(rbt_node_t* node)
{
    return node == NULL ? BLACK : node->color;
}


/**
 * @brief 节点是否是黑色
 * 
 * @param node 
 * @return int 
 */
int is_black(rbt_node_t* node)
{
    return get_node_color(node) == BLACK;
}


/**
 * @brief 节点是否是红色
 * 
 * @param node 
 * @return int 
 */
int is_red(rbt_node_t* node)
{
    return get_node_color(node) == RED;
}


/**
 * @brief 设置节点颜色
 * 
 * @param node 
 * @param color 
 */
rbt_node_t* set_node_color(rbt_node_t* node, int color)
{
    if (node == NULL) {
        return NULL;
    }
    node->color = color;
    return node;
}


/**
 * @brief 设置节点为黑色
 * 
 * @param node 
 */
rbt_node_t* set_black(rbt_node_t* node)
{
    return set_node_color(node, BLACK);
}


/**
 * @brief 设置节点为红色
 * 
 * @param node 
 */
rbt_node_t* set_red(rbt_node_t* node)
{
    return set_node_color(node, RED);
}


/**
 * @brief 获取结点的兄弟节点
 * 
 * @param node 
 * @return rbt_node_t* 
 */
rbt_node_t* get_sibling_ndoe(rbt_node_t* node)
{
    rbt_node_t* parent;

    if (node == NULL) {
        return NULL;
    }

    parent = node->parent;
    if (parent == NULL) {
        return NULL;    // root node
    } else {
        if (is_left_child(node)) {
            return parent->right;
        } else {
            return parent->left;
        }
    }
}
/*--------------------------以上是一些辅助函数--------------------------*/


/**
 * @brief 创建红黑树
 * 
 * @return rbt_t* 
 */
rbt_t* creat_rbt()
{
    rbt_t* tree = malloc(sizeof(rbt_t));
    if (tree == NULL) {
        return NULL;
    }
    return memset(tree, 0, sizeof(rbt_t));
}


/**
 * @brief 销毁红黑树
 * 
 * @param tree 
 * @return int 
 */
int destroy_rbt(rbt_t* tree)
{
    rbt_node_t* stack[10000];
    int top = -1;
    rbt_node_t *cur;

    stack[++top] = tree->root;
    cur = tree->root->left;

    while (cur != NULL || top != -1 ) {
        while (cur != NULL) {
            stack[++top] = cur;
            cur = cur->left;
        }
        
        cur = stack[top]->right;

        free(stack[top--]);
        tree->node_count--;
    }

    free(tree);
    return 0;
}


/**
 * @brief 计算节点是否是叶子节点
 * 
 * @param node 
 * @return 是返回1 不是返回0
 */
int is_leaf_node(rbt_node_t* node)
{
    return (node != NULL) && (node->left == NULL && node->right == NULL);
}


/**
 * @brief 创建节点
 * 
 * @param key: 键值
 * @return 新节点的指针
 */
rbt_node_t* creat_node(int key)
{
    rbt_node_t* new_node;

    new_node = malloc(sizeof(rbt_node_t));
    if (new_node == NULL) {
        return NULL;
    }
    memset(new_node, 0, sizeof(rbt_node_t));

    new_node->val = key;
    return set_red(new_node);
}


/*! 寻找目标节点，找到返回节点，没找到返回NULL */
/**
 * @brief   寻找目标节点，找到返回节点，没找到返回NULL
 * 
 * @param   root: 从此节点开始找
 * @param   key: 键值
 * @return  找到了返回目标节点，没找到返回NULL
 */
rbt_node_t* search_node(rbt_node_t* root, int key)
{
    rbt_node_t* target = root;

    if(root == NULL) {
        return NULL;
    }

    while (target != NULL) {
        if (key < target->val)
        {
            target = target->left;
        } else if (key > target->val) {
            target = target->right;
        } else {
            return target;
        }
    }

    /* 没有目标节点 */
    return NULL;
}


/**
 * @brief 给树上插入一个新节点
 * 
 * @param tree 树
 * @param key 键值
 * @return 返回树本身
 */
rbt_t* insert_node(rbt_t* tree, int key)
{
    rbt_node_t* cur = NULL;
    rbt_node_t* new_node = NULL;


    new_node = creat_node(key);
    if (new_node == NULL) {
        return tree;
    }

    // 如果是空树，则创建根结点
    if (tree->root == NULL) {
        tree->root = new_node;
        tree->node_count++;
        return tree;
    }
    
    cur = tree->root;
    new_node->parent = cur;
    while (cur != NULL) {
        new_node->parent = cur;
        if (key < cur->val) {
            if (cur->left == NULL) {
                cur->left = new_node;
                break;
            } else {
                cur = cur->left;
            }
        } else if (key > cur->val) {
            if (cur->right == NULL) {
                cur->right = new_node;
                break;
            } else {
                cur = cur->right;
            }
        } else {
            // 暂时先不做任何动作
            free(new_node);
            return tree;
        }
    }
    tree->node_count++;

    // 添加完节点，调整平衡
    after_add_node(tree, new_node);
    
    return tree;
}


/**
 * @brief 删除非根节点的叶子节点
 * 
 * @param tree 树
 * @param del 要删除的节点
 * @param flag 要删除的节点是其父节点的左节点还是右节点
 * @return int 
 */
int delete_leaf_node(rbt_t* tree, rbt_node_t* del, int flag)
{
    rbt_node_t* del_parent = del->parent;

    /*! del == root，即要删除的叶子节点是根节点 */
    if (del == tree->root) {
        tree->root = NULL;
    } else {
        if (flag == LEFT_NODE) {
            del_parent->left = NULL;
        }
        if (flag == RIGHT_NODE) {
            del_parent->right = NULL;
        }
    }

    free(del);
    tree->node_count--;
    return 0;
}


/**
 * @brief Get the rsb min node object
 * 
 * @param root 
 * @param parentNode 
 * @return rbt_node_t* 
 */
rbt_node_t* get_rsb_min_node(rbt_node_t* root, rbt_node_t** parentNode)
{
    rbt_node_t* cur;

    if (root == NULL) {
        *parentNode = NULL;
        return NULL;
    }

    *parentNode = root;
    cur = root->right;
    while (cur->left != NULL) {
        *parentNode = cur;
        cur = cur->left;
    }

    return cur;
}


/**
 * @brief 删除键值对应的节点
 * 
 * @param tree 
 * @param key 
 * @return int 
 */
int delete_node(rbt_t* tree, int key)
{
    rbt_node_t *del = NULL, *del_parent = NULL;
    rbt_node_t *rsb_min_node, *rsb_min_node_parent;
    rbt_node_t* replace;
    int flag;


    if (tree == NULL || tree->root == NULL) {
        return 0;
    }

    del = search_node(tree->root, key);
    if (del == NULL) {
        return 0;
    }
    /*! 根结点的父节点是空NULL */
    del_parent = del->parent;

    /*!< 叶子节点 */
    if (del->left == NULL && del->right == NULL) {
        /*! 是根节点 */
        if (del == tree->root) {
            delete_leaf_node(tree, del, -1);
            return 0;
        } else { /*! 不是根节点 */
            /*! 计算节点是其父节点的左节点还是右节点 */
            flag = key < del_parent->val ? LEFT_NODE : RIGHT_NODE;
            delete_leaf_node(tree, del, flag);

            // 删除节点后可能导致失衡
            after_delete_node(tree, del, NULL);
            return 0;
        }
    } else if (del->left == NULL && del->right != NULL) { /*!< 删除节点的左孩子为空，右孩子不为空 */
        if (del == tree->root) { /*! 根结点 */
            tree->root = tree->root->right;
            tree->root->parent = NULL;
        } else {
            /*! 计算节点是其父节点的左节点还是右节点 */
            flag = key < del_parent->val ? LEFT_NODE : RIGHT_NODE;
            if (flag == LEFT_NODE) {
                del->right->parent = del->parent;
                del_parent->left = del->right;        
            } else {
                del->right->parent = del->parent;
                del_parent->right = del->right;
            }
        }
        free(del);
        tree->node_count--;

        replace = del_parent->left == NULL ? del_parent->right : del_parent->right;
        // 删除节点后可能导致失衡
        after_delete_node(tree, del, replace);
    } else if (del->left != NULL && del->right == NULL) {/*!< 删除节点的右孩子为空，左孩子不为空 */
        if (del == tree->root) /*! 根结点 */
        {
            tree->root = tree->root->left;
            tree->root->parent = NULL;
        } else {
            /*! 计算节点是其父节点的左节点还是右节点 */
            flag = key < del_parent->val ? LEFT_NODE : RIGHT_NODE;
            if (flag == LEFT_NODE) {
                del->left->parent = del->parent;
                del_parent->left = del->left;
            } else {
                del->left->parent = del->parent;
                del_parent->right = del->left;
            }
        }
        free(del);
        tree->node_count--;

        replace = del_parent->left == NULL ? del_parent->right : del_parent->right;
        // 删除节点后可能导致失衡
        after_delete_node(tree, del, replace);
    } else { /*!< 删除节点的左右孩子都不为空 */
        rsb_min_node = get_rsb_min_node(del, &rsb_min_node_parent);
        if (del == tree->root) { /*! 根结点 */
            if (del->right->left == NULL) {
                rsb_min_node->left = tree->root->left;
                tree->root->left->parent = rsb_min_node;

                tree->root = rsb_min_node;
                tree->root->parent = NULL;
            } else {
                rsb_min_node_parent->left = rsb_min_node->right;
                if (rsb_min_node_parent->left == NULL) {
                    rsb_min_node->left = tree->root->left;
                    tree->root->left->parent = rsb_min_node;

                    rsb_min_node->right = tree->root->right;
                    tree->root->right->parent = rsb_min_node;

                    tree->root = rsb_min_node;
                    tree->root->parent = NULL;
                } else {
                    rsb_min_node_parent->left->parent = rsb_min_node_parent;

                    rsb_min_node->left = tree->root->left;
                    tree->root->left->parent = rsb_min_node;

                    rsb_min_node->right = tree->root->right;
                    tree->root->right->parent = rsb_min_node;

                    tree->root = rsb_min_node;
                    tree->root->parent = NULL;
                }
            }
            free(del);
            tree->node_count--;
            return 0;
        } else { /*!< 非根节点 */
            flag = key < del_parent->val ? LEFT_NODE : RIGHT_NODE;
            if (flag == LEFT_NODE) { /*! 要删除的节点是其父节点的左节点 */
                if (del == rsb_min_node_parent) {
                    rsb_min_node->parent = del->parent;
                    del->parent->left = rsb_min_node;

                    rsb_min_node->left = del->left;
                    del->left->parent = rsb_min_node;
                } else {
                    rsb_min_node_parent->left = rsb_min_node->right;
                    if (rsb_min_node_parent->left == NULL) {
                        rsb_min_node->left = del->left;
                        del->left->parent = rsb_min_node;

                        rsb_min_node->right = del->right;
                        del->right->parent = rsb_min_node;

                        del->parent->left = rsb_min_node;
                        rsb_min_node->parent = del->parent;
                    } else {
                        rsb_min_node_parent->left->parent = rsb_min_node_parent;

                        rsb_min_node->left = del->left;
                        del->left->parent = rsb_min_node;

                        rsb_min_node->right = del->right;
                        del->right->parent = rsb_min_node;

                        del->parent->left = rsb_min_node;
                        rsb_min_node->parent = del->parent;
                    }
                }
            }
            else { /*! 要删除的节点是其父节点的右节点 */
                if (del == rsb_min_node_parent) {
                    rsb_min_node->parent = del->parent;
                    del->parent->right = rsb_min_node;

                    rsb_min_node->left = del->left;
                    del->left->parent = rsb_min_node;
                } else {
                    rsb_min_node_parent->left = rsb_min_node->right;
                    if (rsb_min_node_parent->left == NULL) {
                        rsb_min_node->left = del->left;
                        del->left->parent = rsb_min_node;

                        rsb_min_node->right = del->right;
                        del->right->parent = rsb_min_node;

                        del->parent->right = rsb_min_node;
                        rsb_min_node->parent = del->parent;
                    } else {
                        rsb_min_node_parent->left->parent = rsb_min_node_parent;

                        rsb_min_node->left = del->left;
                        del->left->parent = rsb_min_node;

                        rsb_min_node->right = del->right;
                        del->right->parent = rsb_min_node;

                        del->parent->right = rsb_min_node;
                        rsb_min_node->parent = del->parent;
                    }
                }
            }
            free(del);
            tree->node_count--;

            replace = del_parent->left == NULL ? del_parent->right : del_parent->right;
            // 删除节点后可能导致失衡
            after_delete_node(tree, del, replace);
            return 0;
        }
    }

    return 0;
}


/**
 * @brief 计算两int型数较大值
 * 
 * @param a 
 * @param b 
 * @return int 
 */
int get_max_val(int a, int b)
{
    return a > b ? a : b;
}


/**
 * @brief 是否是左孩子
 * 
 * @param node 
 * @return int 
 */
int is_left_child(rbt_node_t* node)
{
    if(node == NULL) {
        return 0;
    }
    return node->parent != NULL && node == node->parent->left;
}


/**
 * @brief 是否是右孩子
 * 
 * @param node 
 * @return int 
 */
int is_right_child(rbt_node_t* node)
{
    if (node == NULL) {
        return 0;
    }
    return node->parent != NULL && node == node->parent->right;
}


/**
 * @brief 左旋右旋后半部分步骤相同，单独封装函数
 * 
 * @param node 
 */
static void after_rotate(rbt_t* tree, rbt_node_t* grandpa, rbt_node_t* parent, rbt_node_t* child)
{
    parent->parent = grandpa->parent;
    if (is_left_child(grandpa)) {
        grandpa->parent->left = parent;
    } else if (is_right_child(grandpa)) {
        grandpa->parent->right = parent;
    } else {
        tree->root = parent;
        tree->root->parent = NULL;
    }

    if (child != NULL) {
        child->parent = grandpa;
    }

    grandpa->parent = parent;
}


/**
 * @brief 将此节点左旋
 * 
 * @param grandpa_node 
 * @param parent 
 * @return int 
 */
int rotate_left(rbt_t* tree, rbt_node_t* grandpa)
{
    rbt_node_t *parent, *child;

    if (grandpa == NULL) {
        return -1;
    }

    parent = grandpa->right;
    if (parent == NULL) {
        return -1;
    }

    child = parent->left;
    grandpa->right = child;
    parent->left = grandpa;

    after_rotate(tree, grandpa, parent, child);

    return 0;
}


/**
 * @brief 将此节点右旋
 * 
 * @param grandpa_node 
 * @param parent 
 * @return int 
 */
int rotate_right(rbt_t* tree, rbt_node_t* grandpa)
{
    rbt_node_t *parent, *child;

    if (grandpa == NULL) {
        return -1;
    }


    parent = grandpa->left;
    if (parent == NULL) {
        return -1;
    }

    child = parent->right;
    grandpa->left = child;
    parent->right = grandpa;

    after_rotate(tree, grandpa, parent, child);

    return 0;
}


/**
 * @brief 处理添加节点后导致失衡(新添加节点共12种情况)
 * 
 * @param tree 
 * @param node 新添加的节点
 * @return int 
 */
int after_add_node(rbt_t* tree, rbt_node_t* node)
{
    rbt_node_t *parent, *uncle, *grandpa;

    if (tree == NULL || node == NULL) {
        return -1;
    }

    parent = node->parent;
    uncle = get_sibling_ndoe(parent);

    // 1. 新添加的节点是根节点
    if (parent == NULL) {
        set_black(node); // 把根节点染成黑色
        return 0;
    }
    
    // 2. 新添加的节点的父节点为黑色，不用做调整(4种情况)
    else if (is_black(parent)) {
        return 0;
    }

    // 3. 新添加的节点的叔父节点是红色，上溢(4种情况)
    else if (is_red(uncle)) {
        // 3.1 父节点和叔父节点染成黑色
        set_black(parent);
        set_black(uncle);

        // 3.2 祖父节点向上合并(把祖父节点当作是新添加的节点，新添加的节点默认红色)
        grandpa = parent->parent; // 找到祖父节点
        after_add_node(tree, set_red(grandpa));
        return 0;
    }

    // 4. 新添加的节点的叔父节点不是红色，旋转(4种情况)
    else if (!is_red(uncle)) {
        grandpa = parent->parent; // 找到祖父节点
        if (is_left_child(parent)) { // L
            set_red(grandpa);
            if (is_left_child(node)) { // LL
                set_black(parent);
            } else { // LR
                set_black(node);
                rotate_left(tree, parent);
            }
            rotate_right(tree, grandpa);
        } else { // R
            set_red(grandpa);
            if (is_left_child(node)) { // RL
                set_black(node);
                rotate_right(tree, parent);
            } else { // RR
                set_black(parent);
            }
            rotate_left(tree, grandpa);
        }
        return 0;
    }
    // end
    return 0;
}


/**
 * @brief 处理删除节点后导致失衡
 * 
 * @param tree 
 * @param node 要删除的节点
 * @return int 
 */
int after_delete_node(rbt_t* tree, rbt_node_t* del, rbt_node_t* replace)
{
    rbt_node_t *sibling, *parent;
    int left;
    //int parent_color;

    if (tree == NULL || del == NULL) {
        return -1;
    }

    // 1. 要删除的节点是红色，不用做处理
    if (is_red(del)) {
        return 0;
    }

    // 2. 取代node的节点是红色
    else if (is_red(replace)) {
        set_black(replace);
        return 0;
    }

    // 3. 删除的是黑色叶子节点
    else if (is_black(del)) {
        // 判断被删除的节点是左还是右，哪边被删掉了节点就是哪边的
        parent = del->parent;
        left = parent->left == NULL || is_leaf_node(parent);
        // 拿到要删除结点的兄弟节点
        sibling = left ? parent->right : parent->left;

        // 判断被删除的节点在哪边
        if (left) { // 被删除的节点在左边，兄弟节点在右边
            // 3.1.1 兄弟节点是红色
            if (is_red(sibling)) {
                set_black(sibling);
                set_red(parent);
                rotate_left(tree, parent);

                // 更新兄弟节点
                sibling = parent->right;
            }

            // 3.1.2 兄弟节点是黑色
            if (is_black(sibling->left) && is_black(sibling->right)) { // 兄弟节点没有一个红色节点
                //parent_color = get_node_color(parent); // 保存父节点的颜色
                set_black(parent);
                set_red(sibling);
                if (is_black(parent)) {
                    after_delete_node(tree, parent, NULL);
                }
            } else { // 兄弟节点至少有一个红色节点，向兄弟借一个节点
                // 兄弟节点左边是黑色，兄弟先旋转
                if (is_black(sibling->right)) {
                    rotate_right(tree, sibling);

                    // 更新兄弟节点
                    sibling = parent->right;
                }
                
                // 先染色，再旋转
                set_node_color(sibling, get_node_color(parent));
                set_black(sibling->right);
                set_black(parent);

                // 统一右旋操作
                rotate_left(tree, parent);
            }
        } else { // 被删除的节点在右边，兄弟节点在左边
            // 3.2.1 兄弟节点是红色
            if (is_red(sibling)) {
                set_black(sibling);
                set_red(parent);
                rotate_right(tree, parent);
                sibling = parent->left;
            }

            // 3.2.2 兄弟节点是黑色
            if (is_black(sibling->left) && is_black(sibling->right)) { // 兄弟节点没有一个红色节点
                //parent_color = get_node_color(parent); // 保存父节点的颜色
                set_black(parent);
                set_red(sibling);
                if (is_black(parent)) {
                    after_delete_node(tree, parent, NULL);
                }
            } else { // 兄弟节点至少有一个红色节点，向兄弟借一个节点
                // 兄弟节点左边是黑色，兄弟先旋转
                if (is_black(sibling->left)) {
                    rotate_left(tree, sibling);

                    // 更新兄弟节点
                    sibling = parent->left;
                }
                
                // 先染色，再旋转
                set_node_color(sibling, get_node_color(parent));
                set_black(sibling->left);
                set_black(parent);

                // 统一右旋操作
                rotate_right(tree, parent);
            }
        }
    }

    return 0;
}


/**
 * @brief 中序遍历
 * 
 * @param root 
 */
void infix(rbt_node_t* root)
{
    if (root == NULL) {
        return;
    }

    infix(root->left);
    printf("%d -> %s\n", root->val, root->color == BLACK ? "BLACK" : "RED");
    infix(root->right);
}


/**
 * @brief 层序遍历
 * 
 * @param root 
 * @return int 
 */
int level(rbt_node_t* root)
{
    if (root == NULL) {
        return 0;
    }

    rbt_node_t* queue[10000];
    rbt_node_t* tmp;
    int head = 0;
    int tail = 0;

    queue[tail++] = root;

    while (head != tail) {
        tmp = queue[head++];
        printf("%d ", tmp->val);

        if (tmp->left != NULL) {
            queue[tail++] = tmp->left;
        }

        if (tmp->right != NULL) {
            queue[tail++] = tmp->right;
        }
    }

    return 0;
}






/*********************** END OF FILE ***********************/

